// Date:   Sun Aug 24 10:47:45 AM 2025
// Mail:   lunar_ubuntu@qq.com
// Author: https://github.com/xiaoqixian

#pragma once

#include <any>
#include <memory>
#include <mutex>
#include <string>
#include <unordered_map>
#include <optional>

struct ContextData {
  std::unordered_map<std::string, std::any> kv_;
  mutable std::mutex mtx_;
  std::shared_ptr<ContextData> parent_ = nullptr;
};

class Context {
public:
  using Ptr = std::shared_ptr<ContextData>;

  Context() : data_(std::make_shared<ContextData>()) {}
  Context(Ptr data): data_(data) {}

  Context with_child() {
    auto child = std::make_shared<ContextData>();
    child->parent_ = data_;
    return Context(child);
  }

  template<typename T>
  void set(const std::string& key, T&& value) {
    if (data_ == nullptr) {
      throw std::runtime_error("Empty context");
    }
    std::lock_guard<std::mutex> lk(data_->mtx_);
    data_->kv_[key] = std::forward<T>(value);
  }

  template<typename T>
  std::optional<T> get(const std::string& key) const {
    if (data_ == nullptr) return std::nullopt;
    {
      std::lock_guard<std::mutex> lk(data_->mtx_);
      auto it = data_->kv_.find(key);
      if (it != data_->kv_.end()) {
        try {
          return std::any_cast<T>(it->second);
        } catch (const std::bad_any_cast&) {
          throw std::runtime_error("Type mismatch for key: " + key);
        }
      }
    }
    if (data_->parent_ != nullptr) {
      return Context(data_->parent_).get<T>(key);
    }
    return std::nullopt;
  }

private:
  std::shared_ptr<ContextData> data_;
};
